package goleetcode

import "sort"

func resultArray(nums []int) []int {
	n := len(nums)
	sortNums := make([]int, n)
	copy(sortNums, nums)
	sort.Ints(sortNums)
	index := make(map[int]int)
	for i, num := range sortNums {
		index[num] = i + 1
	}
	arr1, arr2 := []int{nums[0]}, []int{nums[1]}
    tree1, tree2 := NewBinaryIndexedTree(n), NewBinaryIndexedTree(n)
    tree1.Add(index[nums[0]])
    tree2.Add(index[nums[1]])

    for i := 2; i < n; i++ {
        count1 := len(arr1) - tree1.Get(index[nums[i]])
        count2 := len(arr2) - tree2.Get(index[nums[i]])
        if count1 > count2 || (count1 == count2 && len(arr1) <= len(arr2)) {
            arr1 = append(arr1, nums[i])
            tree1.Add(index[nums[i]])
        } else {
            arr2 = append(arr2, nums[i])
            tree2.Add(index[nums[i]])
        }
    }

    return append(arr1, arr2...)
}

type BinaryIndexedTree struct {
	tree []int
}

func NewBinaryIndexedTree(n int) *BinaryIndexedTree {
	return &BinaryIndexedTree{
		tree: make([]int, n+1),
	}
}

func (bt *BinaryIndexedTree) Add(i int) {
	for i < len(bt.tree) {
		bt.tree[i]++
		i += i & -i
	}
}

func (bt *BinaryIndexedTree) Get(i int) int {
	sum := 0
	for i > 0 {
		sum += bt.tree[i]
		i -= i & -i
	}
	return sum
}